package com.computer.fundamentals.datastructure;

import com.computer.fundamentals.datastructure.node.TrieTreeNode;
import com.computer.util.Constant;
import com.computer.util.UniversalMethod;

import java.util.*;

public class TrieTree {

    TrieTreeNode root;

    public TrieTree(TrieTreeNode root) {
        this.root = root;
    }

    /**
     * 添加元素
     */
    public void add(String word) {
        if (isWordContains(word)) {
            throw new RuntimeException(Constant.STRING_IS_ALREADY_EXIST);
        }
        TrieTreeNode cur = root;
        for (Character character: word.toCharArray()) {
            char space = 0;
            int idx = character - space;
            if (cur.count[idx] == 0) {
                cur.children[idx] = new TrieTreeNode();
            }
            cur.count[idx]++;
            cur = cur.children[idx];
        }
        cur.isEnd = true;
    }

    /**
     * 查找word是否在字典树内
     */
    public boolean isWordContains(String word) {
        TrieTreeNode cur = root;
        for (Character character: word.toCharArray()) {
            char space = 0;
            int idx = character - space;
            if (cur.children[idx] != null && cur.count[idx] > 0) {
                cur = cur.children[idx];
            }else {
                return false;
            }
        }
        return cur.isEnd;
    }

    /**
     * 查找prefix是否在字典树内
     */
    public boolean isPrefixContains(String prefix) {
        TrieTreeNode cur = root;
        for (Character character: prefix.toCharArray()) {
            char space = 0;
            int idx = character - space;
            if (cur.children[idx] != null && cur.count[idx] > 0) {
                cur = cur.children[idx];
            }else {
                return false;
            }
        }
        return true;
    }

    /**
     * 删除元素
     */
    public void remove(String word) {
        if (!isWordContains(word)) {
            throw new RuntimeException(Constant.STRING_IS_NOT_EXIST);
        }
        TrieTreeNode cur = root;
        for (Character character: word.toCharArray()) {
            char space = 0;
            int idx = character - space;
            if (cur.count[idx] == 0) {
                throw new RuntimeException(Constant.STRING_LENGTH_ZERO);
            }
            cur.count[idx]--;
        }
    }

    /**
     * 打印字典树
     */
    public void printTrieTree() {
        Queue<TrieTreeNode> queue = new LinkedList<>();
        queue.offer(root);
        while (!queue.isEmpty()) {
            int n = queue.size();
            for (int i = 0;i < n;i++) {
                TrieTreeNode node = queue.poll();
                if (node == null) {
                    continue;
                }
                for (int j = 0;j < Constant.DEFAULT_TRIE_TREE_NODE_SIZE;j++) {
                    if (node.children[j] != null) {
                        System.out.print((char) j);
                        queue.addAll(Arrays.asList(node.children[j]));
                    }
                }
            }
            System.out.print("\n");
        }
    }

    /**
     * 测试
     */
    public static void main(String[] args) {
        TrieTree trieTree = UniversalMethod.createTrieTree();

        System.out.println("---------------完整字典树---------------");
        trieTree.printTrieTree();

        System.out.println("---------------添加元素---------------");
        trieTree.add("father");
        trieTree.printTrieTree();

        System.out.println("---------------检查单词---------------");
        boolean love = trieTree.isWordContains("love");
        System.out.println(love);
        System.out.println("\n");

        System.out.println("---------------检查前缀---------------");
        boolean lov = trieTree.isPrefixContains("lov");
        System.out.println(lov);
        System.out.println("\n");

    }
}
